## ----------------------------- a naive guide for W matrix
## 1. check sparisity: row count and column count of neighbors with positive weights
## 2. check eigen value: the minimum negative eigenvalue should be around -1, 
##                       if the minimum negative eigenvalue is around 0, then the W matrix may not work
## 3. w matrix with complex eigen value (asymmetric w matrix): no worries, just 
##                                                             check 2 for all real eigenvalues



## ----------------------------- output

## [1] "beta"     "rhoy"     "omega"    "sigma2"   "Alpha"   
## [6] "Xi"       "rho0"     "Rho"      "kappa"    "sigma_n2"
## [11] "rhoA"     "L"        "Factor"  

## constant coef
## ar1 parameter for outcome 
## omega for factor loading variance
## error variance
## unit-random coef
## time-random coef
## constant sar coef
## random sar coef
## ar1 parameter for rho_t
## error variance in state space rho_t
## coef for time-vary covar that explains rho_t
## factor loadings matrix 
## factor matrix 

## ---------------------- sim1 block diagnoal w matrix

rm(list = ls())
## compile functions 
library(bpNet)
library(mcmcr)
library(coda)
library(MASS)
library(abind)
setwd('~/Desktop/PSRMreplicationLP/')

set.seed(123456789)

## ------------------------- ##
source("code/simulated20529.R")
source("code/net_plot.R")


## ------------- gen W matrix ------------ ## 

N <- 200 
Ng <- 10 
TT <- 50 

W0 <- matrix(0, N, N)

for (i in 1:20) {
	W0[((i - 1) * 10 + 1) : (i * 10), ((i - 1) * 10 + 1) : (i * 10)] <- 1
}

diag(W0) <- 0

## factor term 
r <- 2 

F <- matrix(rnorm(2 * TT), TT, 2)
L <- matrix(rnorm(2 * N), N, 2)

FL <- F %*% t(L)  ## TT * N

W <- array(NA, dim = c(N, N, TT))

W.new <- W.old <- W0

for (i in 1:TT) {
	## W.new <- W.old 

	## factor term at time i
	subfl <- FL[i, ]

	## choose 4 units and allow them to take an action 
	n <- sample(1:N, 4, replace = FALSE)

	## action 
	na <- rbinom(4, 1, 0.5) ## 0 drop; 1 add

	for (j in 1:4) {
		subw <- W.old[n[j], ]

		
		pos1 <- which(subw == 1)

		pos2 <- which(subw == 0)
		pos2 <- setdiff(pos2, n[j]) ## diagnal is always 0

		if (na[j] == 0) {
			subprob <- exp(abs(subfl[pos1] - subfl[n[j]]))
			sdrop <- which(c(rmultinom(1, 1, subprob)) == 1)
			W.new[n[j], pos1[sdrop]] <- 0

		} else {
			subprob <- exp(-abs(subfl[pos2] - subfl[n[j]]))
			sadd <- which(c(rmultinom(1, 1, subprob)) == 1)
			W.new[n[j], pos2[sadd]] <- 1
		}
	}

	W[,, i] <- W.new
	W.old <- W.new

}

## normalize row sum 

for (i in 1:TT) {
	for (j in 1:N) {
		subsum <- sum(W[j, , i])
		if (subsum > 0) {
			W[j, , i] <- W[j, , i] / subsum
		}
	}
}



## time-varying covar that explains rho_t
rhoZ <- cbind(1, rnorm(TT))


## ------------- other parameters ------------ ## 

## covar
p <- 2
beta <- as.matrix(c(5, 1, 3, -2, 2))

## rho_t
rhoA <- as.matrix(c(0.3, -0.2))
kappa <- 0.3
sigma_n2 <- 0.36

## factor numbers
r <- 2

## error variance
sigma2 <- 1

## contextual effects
contextual.effect <- "both"
force <- "both"


sim <- spatial.simulate(W = W, ## W matrix
                        beta = beta,
		                p = p, 
		                r = 2,
		                contextual.effect = contextual.effect,
		                force = force,
		                sigma2 = sigma2,
		                kappa = kappa, 
		                rhoZ = rhoZ,
		                rhoA = rhoA, 
		                sigma_n2 = sigma_n2, 
		                corH = 0)

data1 <- sim$data
## add a new column specifying group indicator 
## data$group <- c(t(group)) ## group is N * TT while data is sorted by id, time 

## true rho
Rho1 <- sim$Rho


save(W, F, L, data1, Rho1, rhoZ, file = "data/SimData/saom_sim1.RData")

